from torch import nn
from nets.graph_encoder import GraphAttentionEncoder

class CriticNetwork(nn.Module):

    def __init__(
        self,
        input_dim,#输入特征的维度
        embedding_dim,#图节点嵌入的维度
        hidden_dim,#隐藏层的维度
        n_layers,#编码器中图注意力层的数量
        encoder_normalization#编码器中使用的归一化方式
    ):
        super(CriticNetwork, self).__init__()

        self.hidden_dim = hidden_dim

        #该对象用于将输入的节点特征图进行编码，以产生图的嵌入表示。编码器的参数包括节点维度 
        self.encoder = GraphAttentionEncoder(
            node_dim=input_dim,#节点维度
            n_heads=8,#注意力头数量
            embed_dim=embedding_dim,#嵌入维度
            n_layers=n_layers,#注意力层的数量 
            normalization=encoder_normalization#归一化
        )

        self.value_head = nn.Sequential(
            nn.Linear(embedding_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )#两个linear层直接将embedding后的值纬度转为1

    def forward(self, inputs):
        """
        在 forward 方法中，输入张量首先被传递给编码器，得到的输出包括编码后的图结构以及对应的图嵌入表示。
        然后，图嵌入表示被传递给值头模块，经过线性映射和 ReLU 激活函数后，
        最终得到一个单一的值输出，代表了对输入图的值函数估计。
        :param inputs: (batch_size, graph_size, input_dim)
        :return:
        """
        _, graph_embeddings = self.encoder(inputs)
        return self.value_head(graph_embeddings)
